import torch
import torchvision
from torch.utils import data

from 实验.实验11 import Cutout


def getStat(train_data):
    '''
    Compute mean and variance for training data
    :param train_data: 自定义类Dataset(或ImageFolder即可)
    :return: (mean, std)
    '''
    print('Compute mean and variance for training data and testing data.')
    print(len(train_data))
    train_loader = torch.utils.data.DataLoader(
        train_data, batch_size=1, shuffle=False, num_workers=0,
        pin_memory=True)
    mean = torch.zeros(3)
    std = torch.zeros(3)
    for X, _ in train_loader:
        for d in range(3):
            mean[d] += X[:, d, :, :].mean()
            std[d] += X[:, d, :, :].std()
    mean.div_(len(train_data))
    std.div_(len(train_data))
    return list(mean.numpy()), list(std.numpy())


if __name__ == '__main__':
    train_augs = torchvision.transforms.Compose([
        torchvision.transforms.RandomCrop(32, padding=4),  # 对原始 32*32 图像四周各填充4个0像素（40*40），然后随机裁剪成32*32
        torchvision.transforms.RandomHorizontalFlip(),  # 按0.5的概率水平翻转图片
        torchvision.transforms.ToTensor(),  # 将图像从numpy的array转化为pytorch训练需要的tensor
        # 随机选择图像中的一块区域，擦除其像素，主要用来进行数据增强。
        # torchvision.transforms.RandomErasing(p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3)),
        # 使用Cutout擦除其像素进行数据增强。
        # Cutout(n_holes=1, length=16),
    ])

    test_augs = torchvision.transforms.Compose([
        # torchvision.transforms.Resize((32, 32)),
        torchvision.transforms.ToTensor()])

    train_dataset = torchvision.datasets.CIFAR10(root="../data", train=True, transform=train_augs, download=True)
    test_dataset = torchvision.datasets.CIFAR10(root="../data", train=False, transform=test_augs, download=True)

    print('train_dataset:')
    print(getStat(train_dataset))
    print('test_dataset:')
    print(getStat(test_dataset))
